"""
数据预处理和增强
这个脚本用于处理和增强ISIC 2018黑色素瘤数据集，包括重调大小、归一化和数据增强。
"""

import cv2
import numpy as np
import os
import glob
from tqdm import tqdm

# 定义全局参数
HEIGHT = 128  # 目标图像高度
WIDTH = 128   # 目标图像宽度

def load_and_preprocess_image(image_path, mask_path=None, augment=False):
    """
    加载并预处理单个图像及其掩码(如果提供)
    
    参数:
        image_path: 图像文件路径
        mask_path: 掩码文件路径(可选)
        augment: 是否进行数据增强
    返回:
        processed_images: 预处理后的图像列表
        processed_masks: 预处理后的掩码列表(如果mask_path不为None)
    """
    try:
        # 加载图像
        img = cv2.imread(image_path, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"无法加载图像: {image_path}")
        img = cv2.resize(img, (WIDTH, HEIGHT))
        img = img / 255.0  # 归一化
        img = img.astype(np.float32)
        
        processed_images = [img]
        processed_masks = []
        
        # 如果提供了掩码路径
        if mask_path:
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            if mask is None:
                raise ValueError(f"无法加载掩码: {mask_path}")
            mask = cv2.resize(mask, (WIDTH, HEIGHT))
            mask = (mask > 0).astype(np.float32)  # 二值化掩码，使用float32而非int32
            processed_masks = [mask]
        
        # 如果需要数据增强
        if augment and mask_path:
            # 水平翻转 - 简单且可靠的增强方式
            img_h_flip = cv2.flip(img, 1)
            mask_h_flip = cv2.flip(mask, 1)
            processed_images.append(img_h_flip)
            processed_masks.append(mask_h_flip)
            
            # 垂直翻转 - 简单且可靠的增强方式
            img_v_flip = cv2.flip(img, 0)
            mask_v_flip = cv2.flip(mask, 0)
            processed_images.append(img_v_flip)
            processed_masks.append(mask_v_flip)
            
            # 不使用旋转增强，因为它可能会导致错误
        
        if mask_path:
            return processed_images, processed_masks
        else:
            return processed_images
    except Exception as e:
        print(f"处理图像 {image_path} 时出错: {str(e)}")
        raise

def process_dataset(images_dir, masks_dir, output_dir, augment=True):
    """
    处理整个数据集
    
    参数:
        images_dir: 图像目录
        masks_dir: 掩码目录
        output_dir: 输出目录
        augment: 是否进行数据增强
    """
    print(f"处理目录：\n图像: {images_dir}\n掩码: {masks_dir}\n输出: {output_dir}")
    
    # 检查目录是否存在
    if not os.path.exists(images_dir):
        print(f"错误: 图像目录 '{images_dir}' 不存在")
        return None, None
    
    if not os.path.exists(masks_dir):
        print(f"错误: 掩码目录 '{masks_dir}' 不存在")
        return None, None
        
    if not os.path.exists(output_dir):
        print(f"创建输出目录: {output_dir}")
        os.makedirs(output_dir)
    
    # 获取所有图像文件
    image_files = sorted(glob.glob(os.path.join(images_dir, '*.jpg')))
    print(f"找到 {len(image_files)} 个图像文件")
    
    all_images = []
    all_masks = []
    
    # 为每个图像文件查找对应的掩码文件
    processed_count = 0
    for img_file in tqdm(image_files, desc="处理图像"):
        # 获取图像ID (例如 ISIC_0012345)
        img_id = os.path.basename(img_file).split('.')[0]
        # 构造对应的掩码文件名 (例如 ISIC_0012345_segmentation.png)
        mask_file_name = f"{img_id}_segmentation.png"
        mask_file_path = os.path.join(masks_dir, mask_file_name)
        
        # 检查掩码文件是否存在
        if not os.path.exists(mask_file_path):
            print(f"警告: 找不到图像 {img_id} 的对应掩码文件")
            continue
          # 预处理并可选地增强
        try:
            images, masks = load_and_preprocess_image(img_file, mask_file_path, augment)
            all_images.extend(images)
            all_masks.extend(masks)
            processed_count += 1
            if processed_count % 100 == 0:
                print(f"已成功处理 {processed_count} 张图像")
        except Exception as e:
            print(f"处理图像 {img_id} 时出错: {str(e)}")
            continue
    
    # 转换为numpy数组
    all_images_np = np.array(all_images)
    all_masks_np = np.array(all_masks)
    all_masks_np = all_masks_np[..., np.newaxis]  # 添加通道维度
    
    print(f"处理后形状: 图像 {all_images_np.shape}, 掩码 {all_masks_np.shape}")
    
    # 保存处理后的数据
    np.save(os.path.join(output_dir, 'melanoma_images.npy'), all_images_np)
    np.save(os.path.join(output_dir, 'melanoma_masks.npy'), all_masks_np)
    
    print(f"数据已保存至 {output_dir}")
    
    return all_images_np, all_masks_np

def split_dataset(images, masks, train_ratio=0.8):
    """
    将数据集分为训练集和测试集
    
    参数:
        images: 图像数组
        masks: 掩码数组
        train_ratio: 训练集比例
    返回:
        train_images, train_masks, test_images, test_masks: 分割后的数据集
    """
    # 打乱数据顺序
    indices = np.arange(images.shape[0])
    np.random.shuffle(indices)
    
    # 计算训练集大小
    train_size = int(images.shape[0] * train_ratio)
    
    # 分割数据集
    train_indices = indices[:train_size]
    test_indices = indices[train_size:]
    
    train_images = images[train_indices]
    train_masks = masks[train_indices]
    test_images = images[test_indices]
    test_masks = masks[test_indices]
    
    return train_images, train_masks, test_images, test_masks

if __name__ == "__main__":
    # 设置正确的目录路径
    images_dir = "e:/black/data/raw/ISIC2018_Task1-2_Training_Input"
    masks_dir = "e:/black/data/raw/ISIC2018_Task1_Training_GroundTruth"
    output_dir = "e:/black/data/processed"
    
    print("开始处理数据集...")
    try:
        # 启用详细输出
        images, masks = process_dataset(images_dir, masks_dir, output_dir)
        if images is not None and masks is not None:
            # 分割数据集为训练集和测试集
            train_imgs, train_masks, val_imgs, val_masks = split_dataset(images, masks)
            print(f"数据集分割完成: 训练集 {train_imgs.shape[0]} 样本, 验证集 {val_imgs.shape[0]} 样本")
            
            # 保存分割后的数据集
            np.save(os.path.join(output_dir, 'train_images.npy'), train_imgs)
            np.save(os.path.join(output_dir, 'train_masks.npy'), train_masks)
            np.save(os.path.join(output_dir, 'val_images.npy'), val_imgs)
            np.save(os.path.join(output_dir, 'val_masks.npy'), val_masks)
            print("训练和验证数据已保存!")
        else:
            print("数据处理失败，请检查目录配置和文件可用性。")
    except Exception as e:
        print(f"处理过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()
    print("数据处理完成!")
